use std::{sync::{atomic::AtomicU64, Arc}, time::Duration, net::SocketAddr};

use tokio::{sync::{mpsc::{UnboundedSender, UnboundedReceiver}, oneshot::Sender}, net::{TcpStream, tcp::{OwnedReadHalf, OwnedWriteHalf}}, io::{AsyncReadExt, AsyncWriteExt}, time::Instant};

#[derive(Debug)]
pub struct Message(pub Vec<u8>);

pub trait ToMessage {
    fn to_message(&self) -> Vec<u8>;
}

#[derive(Debug)]
enum ConnectionEvent {
    ConnectFailed(u64),
    Connected(u64, SocketAddr, UnboundedSender<Message>, Sender<()>),
    Disconnected(u64),
    Message(u64, Message),
}

async fn read_message(recv_socket: &mut OwnedReadHalf) -> tokio::io::Result<Message> {
    let len = recv_socket.read_u32().await?;
    let mut buffer = Vec::with_capacity(len as usize);
    recv_socket.read_exact(&mut buffer).await?;

    Ok(Message(buffer))
}

async fn write_message(send_socket: &mut OwnedWriteHalf, message: Message) -> tokio::io::Result<()> {
    send_socket.write_u32(message.0.len() as u32).await?;
    send_socket.write_all(&message.0).await?;
    Ok(())
}

pub trait NetEventHandle {
    fn on_connect_failed(&mut self, id: u64);
    fn on_connected(&mut self, id: u64, addr: SocketAddr, message_sender: UnboundedSender<Message>, closer_tx: Sender<()>);
    fn on_disconnected(&mut self, id: u64);
    fn on_message(&mut self, id: u64, message: &[u8]);
}

pub struct MessageQueue {
    tx: UnboundedSender<ConnectionEvent>,
    rx: UnboundedReceiver<ConnectionEvent>,
}

impl MessageQueue {
    pub fn new() -> Self {
        let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<ConnectionEvent>();
        Self {
            tx,
            rx,
        }
    }

    pub async fn run<H: NetEventHandle>(&mut self, handler: &mut H, duration: Duration) {
        let instant = Instant::now() + duration;
        loop {
            tokio::select! {
                _ = tokio::time::sleep_until(instant) => {
                    break;
                }
                Some(event) = self.rx.recv() => {
                    match event {
                        ConnectionEvent::ConnectFailed(id) => {
                            handler.on_connect_failed(id);
                        }
                        ConnectionEvent::Connected(id, addr, message_sender, close_tx) => {
                            handler.on_connected(id, addr, message_sender, close_tx);
                        }
                        ConnectionEvent::Disconnected(id) => {
                            handler.on_disconnected(id);
                        }
                        ConnectionEvent::Message(id, message) => {
                            handler.on_message(id, &message.0);
                        }
                    }
                }
            }
        }
    }
}

pub struct Net {
    next_id: Arc<AtomicU64>,
}

async fn process_new_income_connection(socket: TcpStream, addr: SocketAddr, tx: UnboundedSender<ConnectionEvent>, id: u64) {
    let (mut recv_socket, send_socket) = socket.into_split();

    let (send_queue, recv_queue) = tokio::sync::mpsc::unbounded_channel::<Message>();
    let (closer_send_tx, mut closer_send_rx) = tokio::sync::oneshot::channel::<()>();
    tokio::spawn(async move {
        let mut send_socket = send_socket;
        let mut recv_queue = recv_queue;
        loop {
            tokio::select! {
                _ = &mut closer_send_rx => {
                    break;
                }
                message = tokio::time::timeout(Duration::from_secs(30), recv_queue.recv()) => {
                    let message = match message {
                        Ok(message) => {
                            match message {
                                Some(message) => message,
                                None => {
                                    continue;
                                }
                            }
                        }
                        Err(_) => {
                            Message(Vec::new())
                        }
                    };

                    match write_message(&mut send_socket, message).await {
                        Ok(_) => {}
                        Err(e) => {
                            log::warn!("Error connection: {} writing message: {}", id, e);
                        }
                    }
                }
            }
        }
    });

    let (closer_recv_tx, mut closer_recv_rx) = tokio::sync::oneshot::channel::<()>();
    let (wait_recv_tx, mut wait_recv_rx) = tokio::sync::oneshot::channel::<()>();
    let tx_clone = tx.clone();
    tokio::spawn(async move {
        loop {
            let _ = (&mut wait_recv_rx).await.unwrap();
            tokio::select! {
                _ = &mut closer_recv_rx => {
                    break;
                }
                readed = tokio::time::timeout(Duration::from_secs(30), read_message(&mut recv_socket)) => {
                    match readed {
                        Ok(readed) => {
                            match readed {
                                Ok(message) => {
                                    if message.0.len() > 0 {
                                        tx_clone.send(ConnectionEvent::Message(id, message)).unwrap_or_default();
                                    }
                                }
                                Err(e) => {
                                    log::error!("Error connection: {} reading message: {}", id, e);
                                    break;
                                }
                            }
                        }
                        Err(e) => {
                            log::error!("Error connection: {} reading message: {} time out", id, e);
                            break;
                        }
                    }
                }
            }
        }

        closer_send_tx.send(()).unwrap_or_default();
        tx_clone.send(ConnectionEvent::Disconnected(id)).unwrap();
    });

    tx.send(ConnectionEvent::Connected(id, addr, send_queue, closer_recv_tx)).unwrap();
    wait_recv_tx.send(()).unwrap_or_default();
}

impl Net {
    pub fn new() -> Self {
        Self {
            next_id: Arc::new(AtomicU64::new(1)),
        }
    }

    pub async fn listen(&self, addr: &str, message_queue: Option<MessageQueue>) -> tokio::io::Result<MessageQueue> {
        let listener = tokio::net::TcpListener::bind(addr).await?;
        let message_queue = message_queue.unwrap_or_else(|| MessageQueue::new());
        let next_id = self.next_id.clone();
        let tx = message_queue.tx.clone();
        tokio::spawn(async move {
            loop {
                match listener.accept().await {
                    Ok((socket, addr)) => {
                        let id = next_id.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                        process_new_income_connection(socket, addr, tx.clone(), id).await;
                    }
                    Err(e) => {
                        log::error!("Error listening: {}", e);
                        break;
                    }
                }
            }
        });

        Ok(message_queue)
    }

    pub async fn connect(&self, addr: &str, message_queue: Option<MessageQueue>) -> tokio::io::Result<(u64, MessageQueue)> {
        let addr = addr.to_string();
        let message_queue = message_queue.unwrap_or_else(|| MessageQueue::new());
        let id = self.next_id.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        let tx = message_queue.tx.clone();
        tokio::spawn(async move {
            match tokio::net::TcpStream::connect(&addr).await {
                Ok(socket) => {
                    let addr = socket.peer_addr().unwrap();
                    process_new_income_connection(socket, addr, tx, id).await;
                }
                Err(e) => {
                    log::error!("Error connecting {} : {}", &addr, e);
                    tx.send(ConnectionEvent::ConnectFailed(id)).unwrap_or_default();
                }
            }
        });

        Ok((id, message_queue))
    }

}